# -*- coding: utf-8 -*-
# @Time: 2025/6/8 13:53
# @Author: wzd
# @Email: 2146333089@qq.com
# @File: exact_content.py
"""
    功能：这个函数主要用于对电子档格式的教学书籍进行解析

"""


import os
import re
import pdfplumber
import docx
import pptx
from bs4 import BeautifulSoup
import xml.etree.ElementTree as ET
from typing import List, Dict, Tuple, Optional


class DocumentExtractor:
    """文档内容提取器，支持多种格式文档的概念、公式和大纲提取"""

    def __init__(self):
        # 定义概念提取的正则模式
        self.concept_pattern = re.compile(
            r'(?P<title>[\w\s-]+?)[：:\s]*(?P<description>.*?[。？！;；])',
            re.DOTALL | re.MULTILINE
        )

        # 定义公式提取的正则模式
        self.formula_pattern = re.compile(
            r'(\$.*?\$)|(\\\[.*?\\\])|(<math.*?>.*?</math>)|([A-Za-z0-9_]+\s*=\s*.*?)',
            re.DOTALL
        )

        # 定义大纲提取的正则模式
        self.outline_pattern = re.compile(
            r'^(?P<level>第?[一二三四五六七八九十0-9]{1,6}[章节部分篇卷]?)\s*(?P<title>.*)$',
            re.MULTILINE
        )

        # 定义标题级别检测模式
        self.header_patterns = {
            1: re.compile(r'^第[一二三四五六七八九十0-9]{1,6}章\s*.*$'),
            2: re.compile(r'^第[一二三四五六七八九十0-9]{1,6}节\s*.*$'),
            3: re.compile(r'^[一二三四五六七八九十、]{1,6}\s*.*$'),
            4: re.compile(r'^[1-9][0-9]?[\.、]\s*.*$'),
            5: re.compile(r'^[a-zA-Z]\.\s*.*$'),
            6: re.compile(r'^[ivxlcdm]+\.\s*.*$', re.IGNORECASE)
        }

    def extract_from_pdf(self, file_path: str) -> Dict[str, List[str]]:
        """从PDF文件中提取内容"""
        result = {'concepts': [], 'formulas': [], 'outline': []}

        try:
            with pdfplumber.open(file_path) as pdf:
                full_text = ""
                for page in pdf.pages:
                    page_text = page.extract_text()
                    if page_text:
                        full_text += page_text + "\n"

                # 提取概念
                concepts = self._extract_concepts(full_text)
                result['concepts'] = concepts

                # 提取公式
                formulas = self._extract_formulas(full_text)
                result['formulas'] = formulas

                # 提取大纲
                outline = self._extract_outline(full_text)
                result['outline'] = outline

                # 提取分层标题
                headers = self._extract_headers(full_text)
                result['headers'] = headers

        except Exception as e:
            print(f"Error processing PDF {file_path}: {e}")

        return result

    def extract_from_docx(self, file_path: str) -> Dict[str, List[str]]:
        """从Word文档中提取内容"""
        result = {'concepts': [], 'formulas': [], 'outline': []}

        try:
            doc = docx.Document(file_path)
            full_text = ""

            # 提取文本内容
            for para in doc.paragraphs:
                full_text += para.text + "\n"

            # 提取概念
            concepts = self._extract_concepts(full_text)
            result['concepts'] = concepts

            # 提取公式
            formulas = self._extract_formulas(full_text)
            result['formulas'] = formulas

            # 提取大纲
            outline = self._extract_outline(full_text)
            result['outline'] = outline

            # 提取分层标题
            headers = self._extract_headers(full_text)
            result['headers'] = headers

            # 提取Word中的公式（OLE对象）
            for rel in doc.part.rels.values():
                if "oleObject" in rel.target_ref:
                    # 尝试解析OLE对象中的公式
                    pass

        except Exception as e:
            print(f"Error processing DOCX {file_path}: {e}")

        return result

    def extract_from_pptx(self, file_path: str) -> Dict[str, List[str]]:
        """从PPTX文档中提取内容"""
        result = {'concepts': [], 'formulas': [], 'outline': []}

        try:
            prs = pptx.Presentation(file_path)
            full_text = ""

            # 提取幻灯片文本
            for slide in prs.slides:
                for shape in slide.shapes:
                    if not shape.has_text_frame:
                        continue
                    for paragraph in shape.text_frame.paragraphs:
                        for run in paragraph.runs:
                            full_text += run.text + " "
                        full_text += "\n"
                full_text += "\n"

            # 提取概念
            concepts = self._extract_concepts(full_text)
            result['concepts'] = concepts

            # 提取公式
            formulas = self._extract_formulas(full_text)
            result['formulas'] = formulas

            # 提取大纲（幻灯片标题）
            outline = []
            for i, slide in enumerate(prs.slides):
                title = ""
                if slide.shapes.title:
                    title = slide.shapes.title.text
                    outline.append(f"幻灯片 {i + 1}: {title}")
            result['outline'] = outline

        except Exception as e:
            print(f"Error processing PPTX {file_path}: {e}")

        return result

    def _extract_concepts(self, text: str) -> List[str]:
        """从文本中提取概念"""
        concepts = []
        for match in self.concept_pattern.finditer(text):
            concept = f"{match.group('title').strip()}: {match.group('description').strip()}"
            concepts.append(concept)
        return concepts

    def _extract_formulas(self, text: str) -> List[str]:
        """从文本中提取公式"""
        formulas = []
        for match in self.formula_pattern.finditer(text):
            if match.group(1):  # $...$ 格式
                formulas.append(match.group(1))
            elif match.group(2):  # \[...\] 格式
                formulas.append(match.group(2))
            elif match.group(3):  # <math>...</math> 格式
                # 解析MathML
                try:
                    soup = BeautifulSoup(match.group(3), 'lxml')
                    formulas.append(soup.get_text())
                except:
                    formulas.append(match.group(3))
            elif match.group(4):  # 简单公式格式
                # 检查是否包含等号或数学符号
                if '=' in match.group(4) or '+' in match.group(4) or '-' in match.group(4):
                    formulas.append(match.group(4))
        return formulas

    def _extract_outline(self, text: str) -> List[str]:
        """从文本中提取大纲"""
        outline = []
        for match in self.outline_pattern.finditer(text):
            outline_item = f"{match.group('level')} {match.group('title')}"
            outline.append(outline_item)
        return outline

    def _extract_headers(self, text: str) -> List[Tuple[int, str]]:
        """从文本中提取分层标题"""
        headers = []
        lines = text.split('\n')

        for line in lines:
            line = line.strip()
            if not line:
                continue

            for level, pattern in self.header_patterns.items():
                if pattern.match(line):
                    headers.append((level, line))
                    break  # 一旦匹配到一个级别，就不再检查更低的级别

        return headers

    def save_results(self, results: Dict[str, List[str]], output_dir: str, base_name: str):
        """保存提取的结果到文件"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        for category, items in results.items():
            if not items:
                continue

            output_file = os.path.join(output_dir, f"{base_name}_{category}.txt")
            with open(output_file, 'w', encoding='utf-8') as f:
                if category == 'headers':
                    # 处理分层标题
                    for level, title in items:
                        indent = '  ' * (level - 1)
                        f.write(f"{indent}{title}\n")
                else:
                    # 处理其他类别
                    for item in items:
                        f.write(f"{item}\n")

            print(f"已保存 {category} 到 {output_file}")


def main():
    """主函数，演示如何使用提取器"""
    extractor = DocumentExtractor()

    # # 示例：处理PDF文件
    # pdf_file = "course_material.pdf"  # 替换为实际的文件路径
    # if os.path.exists(pdf_file):
    #     pdf_results = extractor.extract_from_pdf(pdf_file)
    #     extractor.save_results(pdf_results, "output", "pdf")

    # 示例：处理Word文档
    docx_file = r"C:\cnsoft\电子书籍资料\《嵌入式Linux开发实践教程》示例资源-word版课件\cp07-样章示例-TensorFlow.js应用开发.docx"  # 替换为实际的文件路径
    if os.path.exists(docx_file):
        docx_results = extractor.extract_from_docx(docx_file)
        extractor.save_results(docx_results, "output", "docx")
    #
    # # 示例：处理PPTX文档
    # pptx_file = "course_material.pptx"  # 替换为实际的文件路径
    # if os.path.exists(pptx_file):
    #     pptx_results = extractor.extract_from_pptx(pptx_file)
    #     extractor.save_results(pptx_results, "output", "pptx")


if __name__ == "__main__":
    main()